{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwIF4LX30vzc"
      },
      "source": [
        "## Building a Arxiv Research Agent with Pinecone and Llamaindex Workflows\n",
        "\n",
        "Follow along in this notebook to use Pinecone and Llamaindex workflows to learn agentic RAG."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Ak5WKR20ygw"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/learn/generation/llama-index/llama-index-research-agent-condensed.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "WThmiBfvW7o0"
      },
      "outputs": [],
      "source": [
        "!pip install -qU \\\n",
        "    datasets==2.19.1 \\\n",
        "    llama-index-core==0.11.9 \\\n",
        "    llama-index-llms-openai \\\n",
        "    llama-index-utils-workflow==0.2.1 \\\n",
        "    llama-index-embeddings-openai \\\n",
        "    pinecone-notebooks \\\n",
        "    pinecone"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RSNvnnmjhEZx"
      },
      "source": [
        "# Knowledge Base Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yZCzed9thKZm"
      },
      "source": [
        "We'll be running our agent against a knowledge base \u2014 which requires a Pinecone index to be built.\n",
        "\n",
        "You can, if needed, skip this step and replace the `search` tool with a placeholder value if wanting to quickly test the structure of the Llama Index Workflow.\n",
        "\n",
        "If you want full functionality here, you do need to run this section \u2014 but we'll make it quick."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kr1wMhSG09Yr"
      },
      "source": [
        "## Before you begin:\n",
        "\n",
        "You'll need an OpenAI API key and a Pinecone API key to complete this tutorial!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EH0zitbAhI9N"
      },
      "source": [
        "## Download a Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h5UZcVZ1hO-O"
      },
      "source": [
        "The first thing we need for an agent using RAG is somewhere we want to pull knowledge from. We will use v2 of the AI ArXiv dataset, available on Hugging Face Datasets at `jamescalam/ai-arxiv2-semantic-chunks`.\n",
        "\n",
        "Note: we're using the prechunked dataset. For the raw version see `jamescalam/ai-arxiv2`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_LTq8wOWhesx",
        "outputId": "76bb7efc-431b-43f8-e016-361faa1dcdfc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['id', 'title', 'content', 'prechunk_id', 'postchunk_id', 'arxiv_id', 'references'],\n",
              "    num_rows: 209760\n",
              "})"
            ]
          },
          "execution_count": 38,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"jamescalam/ai-arxiv2-semantic-chunks\", split=\"train\")\n",
        "dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mqyP3a0_hqPm",
        "outputId": "c351af9d-caf7-4f65-a5ae-3c414ce72db5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'id': '2401.04088#0',\n",
              " 'title': 'Mixtral of Experts',\n",
              " 'content': '4 2 0 2 n a J 8 ] G L . s c [ 1 v 8 8 0 4 0 . 1 0 4 2 : v i X r a # Mixtral of Experts Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, L\u00c3\u00a9lio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Th\u00c3\u00a9ophile Gervet, Thibaut Lavril, Thomas Wang, Timoth\u00c3\u00a9e Lacroix, William El Sayed Abstract We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine- tuned to follow instructions, Mixtral 8x7B \u00e2 Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B \u00e2 chat model on human bench- marks. Both the base and instruct models are released under the Apache 2.0 license.',\n",
              " 'prechunk_id': '',\n",
              " 'postchunk_id': '2401.04088#1',\n",
              " 'arxiv_id': '2401.04088',\n",
              " 'references': ['1905.07830']}"
            ]
          },
          "execution_count": 39,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dataset[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TO1-zComhw3J"
      },
      "source": [
        "## Construct Knowledge Base"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cztd2YZ8h5VQ"
      },
      "source": [
        "First initialize encoder model, for this we will need an [OpenAI API key](https://platform.openai.com/api-keys)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "ZU92MBJxiUwq"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "from llama_index.embeddings.openai import OpenAIEmbedding\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\") or getpass(\"OpenAI API key: \")\n",
        "\n",
        "# batch size parameter aligns with Pinecone upsertion\n",
        "embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\", embed_batch_size=128)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mpMLbVDLi1pC"
      },
      "source": [
        "Initialize our connection to Pinecone:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 247
        },
        "id": "LzsjnjuCi3iy",
        "outputId": "f38fa6f4-4061-47b6-cef3-6d2c6cba062f"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<script type=\"text/javascript\" src=\"https://connect.pinecone.io/embed.js\"></script>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from pinecone import Pinecone\n",
        "from pinecone_notebooks.colab import Authenticate\n",
        "\n",
        "Authenticate()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "NzolxgoIJ-YZ"
      },
      "outputs": [],
      "source": [
        "# initialize connection to pinecone (get API key at app.pinecone.io)\n",
        "api_key = os.environ.get('PINECONE_API_KEY')\n",
        "\n",
        "# configure client\n",
        "pc = Pinecone(api_key=api_key)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "GS535EvejxFM"
      },
      "outputs": [],
      "source": [
        "from pinecone import ServerlessSpec\n",
        "\n",
        "spec = ServerlessSpec(\n",
        "    cloud=\"aws\", region=\"us-east-1\"  # us-east-1\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gFamwXVekeUx",
        "outputId": "82dfb7a0-c84a-4bd5-cbb8-3dacba20ae1e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1536"
            ]
          },
          "execution_count": 44,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dims = len(embed_model.get_text_embedding(\"some random text\"))\n",
        "dims\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YBzhQFERkjwN",
        "outputId": "259b6917-ccd5-4454-a6cc-4e830821e484"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'dimension': 1536,\n",
              " 'index_fullness': 0.0,\n",
              " 'namespaces': {'': {'vector_count': 10000}},\n",
              " 'total_vector_count': 10000}"
            ]
          },
          "execution_count": 45,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import time\n",
        "\n",
        "index_name = \"gpt-4o-research-agent\"\n",
        "\n",
        "# check if index already exists (it shouldn't if this is first time)\n",
        "if index_name not in pc.list_indexes().names():\n",
        "    # if does not exist, create index\n",
        "    pc.create_index(\n",
        "        index_name,\n",
        "        dimension=dims,  # dimensionality of embed 3\n",
        "        metric='dotproduct',\n",
        "        spec=spec\n",
        "    )\n",
        "    # wait for index to be initialized\n",
        "    while not pc.describe_index(index_name).status['ready']:\n",
        "        time.sleep(1)\n",
        "\n",
        "# connect to index\n",
        "index = pc.Index(index_name)\n",
        "time.sleep(1)\n",
        "# view index stats\n",
        "index.describe_index_stats()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SEP_TzsMkldk"
      },
      "source": [
        "Populate the knowledge base:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "27cb74d36c8247cab8aec50b6fb297ca",
            "e9be8a886ce148c094198d7ed0b7b00d",
            "b4950ebf8a494d7fbed603b6b5273ae4",
            "bb082cb984154fe8aa9bd3d11f6a21b0",
            "3844e9e3b5134fcca305b1ea6cb0c6f4",
            "b4bb0ab8534147eaba3275ed4abd6d9e",
            "3f066832ff284376a5cb01514a4a5384",
            "39d1d92259e94a3484d3c0c06346ac7f",
            "0c9a091794664ca2bd19dfbeb9cab17f",
            "a18107f80dd74c018805efa30af5d547",
            "e8bbb14c0ca04beaa323394578ba3254"
          ]
        },
        "id": "SYdvU-fjkold",
        "outputId": "9dbe8840-4f21-4900-f37b-a0aaafc6b21a"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "27cb74d36c8247cab8aec50b6fb297ca",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/79 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "# easier to work with dataset as pandas dataframe\n",
        "# take first 10k instances as example. Embed more at your convenience\n",
        "data = dataset.to_pandas().iloc[:10000]\n",
        "\n",
        "batch_size = 128\n",
        "\n",
        "for i in tqdm(range(0, len(data), batch_size)):\n",
        "    i_end = min(len(data), i+batch_size)\n",
        "    batch = data[i:i_end].to_dict(orient=\"records\")\n",
        "    # get batch of data\n",
        "    metadata = [{\n",
        "        \"title\": r[\"title\"],\n",
        "        \"content\": r[\"content\"],\n",
        "        \"arxiv_id\": r[\"arxiv_id\"],\n",
        "        \"references\": r[\"references\"].tolist()\n",
        "    } for r in batch]\n",
        "    # generate unique ids for each chunk\n",
        "    ids = [r[\"id\"] for r in batch]\n",
        "    # get text content to embed\n",
        "    content = [r[\"content\"] for r in batch]\n",
        "    # embed text\n",
        "    embeds = embed_model.get_text_embedding_batch(content)\n",
        "    # add to Pinecone\n",
        "    index.upsert(vectors=zip(ids, embeds, metadata))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FUh_C_5HhYh6"
      },
      "source": [
        "# Agent Components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0RS6Ow8FLmZE"
      },
      "source": [
        "## Tools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51qdvTLIYKL6"
      },
      "source": [
        "We define the separate tool functions. When integrating with our graph all of these will be executed using the same `run_tools` class - which we will define later.\n",
        "\n",
        "For now, let's define the functions that our agent will have access to."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "Gxe_NBTHLm01"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "import re\n",
        "\n",
        "# our regex\n",
        "abstract_pattern = re.compile(\n",
        "    r'\\s*Abstract:\\s*(.*?)\\s*',\n",
        "    re.DOTALL\n",
        ")\n",
        "\n",
        "async def fetch_arxiv(arxiv_id: str):\n",
        "    \"\"\"Gets the abstract from an ArXiv paper given the arxiv ID. Useful for\n",
        "    finding high-level context about a specific paper.\"\"\"\n",
        "    print(\">>> fetch_arxiv\")\n",
        "    # get paper page in html\n",
        "    res = requests.get(\n",
        "        f\"https://export.arxiv.org/abs/{arxiv_id}\"\n",
        "    )\n",
        "    # search html for abstract\n",
        "    re_match = abstract_pattern.search(res.text)\n",
        "    # return abstract text\n",
        "    return re_match.group(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NeckltzgL92n"
      },
      "source": [
        "## Rag Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uEU4NDWlL-Ox"
      },
      "source": [
        "We provide two RAG-focused tools for our agent. The `rag_search` allows the agent to perform a simple RAG search for some information across all indexed research papers. The `rag_search_filter` also searches, but within a specific paper which is filtered for via the `arxiv_id` parameter.\n",
        "\n",
        "We also define the `format_rag_contexts` function to handle the transformation of our Pinecone results from a JSON object to a readble plaintext format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "id": "kuv6bK99L-no"
      },
      "outputs": [],
      "source": [
        "def format_rag_contexts(matches: list):\n",
        "    contexts = []\n",
        "    for x in matches:\n",
        "        text = (\n",
        "            f\"Title: {x['metadata']['title']}\\n\"\n",
        "            f\"Content: {x['metadata']['content']}\\n\"\n",
        "            f\"ArXiv ID: {x['metadata']['arxiv_id']}\\n\"\n",
        "            f\"Related Papers: {x['metadata']['references']}\\n\"\n",
        "        )\n",
        "        contexts.append(text)\n",
        "    context_str = \"\\n---\\n\".join(contexts)\n",
        "    return context_str\n",
        "\n",
        "async def rag_search_filter(query: str, arxiv_id: str):\n",
        "    \"\"\"Finds information from our ArXiv database using a natural language query\n",
        "    and a specific ArXiv ID. Allows us to learn more details about a specific paper.\"\"\"\n",
        "    print(\">>> rag_search_filter\")\n",
        "    xq = await embed_model.aget_text_embedding(query)\n",
        "    xc = index.query(vector=xq, top_k=6, include_metadata=True, filter={\"arxiv_id\": arxiv_id})\n",
        "    context_str = format_rag_contexts(xc[\"matches\"])\n",
        "    return context_str\n",
        "\n",
        "async def rag_search(query: str):\n",
        "    \"\"\"Finds specialist information on AI using a natural language query.\"\"\"\n",
        "    print(\">>> rag_search\")\n",
        "    xq = await embed_model.aget_text_embedding(query)\n",
        "    xc = index.query(vector=xq, top_k=2, include_metadata=True)\n",
        "    context_str = format_rag_contexts(xc[\"matches\"])\n",
        "    return context_str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "id": "OGsCy9ZMVJmy"
      },
      "outputs": [],
      "source": [
        "async def final_answer(\n",
        "    introduction: str,\n",
        "    research_steps: str,\n",
        "    main_body: str,\n",
        "    conclusion: str,\n",
        "    sources: str\n",
        "):\n",
        "    \"\"\"Returns a natural language response to the user in the form of a research\n",
        "    report. There are several sections to this report, those are:\n",
        "    - `introduction`: a short paragraph introducing the user's question and the\n",
        "    topic we are researching.\n",
        "    - `research_steps`: a few bullet points explaining the steps that were taken\n",
        "    to research your report.\n",
        "    - `main_body`: this is where the bulk of high quality and concise\n",
        "    information that answers the user's question belongs. It is 3-4 paragraphs\n",
        "    long in length.\n",
        "    - `conclusion`: this is a short single paragraph conclusion providing a\n",
        "    concise but sophisticated view on what was found.\n",
        "    - `sources`: a bulletpoint list provided detailed sources for all information\n",
        "    referenced during the research process\n",
        "    \"\"\"\n",
        "    print(\">>> final_answer\")\n",
        "    if type(research_steps) is list:\n",
        "        research_steps = \"\\n\".join([f\"- {r}\" for r in research_steps])\n",
        "    if type(sources) is list:\n",
        "        sources = \"\\n\".join([f\"- {s}\" for s in sources])\n",
        "    return \"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WUm29WrkVNYh"
      },
      "source": [
        "## Oracle LLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sqM2mkxsVOoG"
      },
      "source": [
        "Our prompt for the Oracle will emphasize it's decision making ability within the `system_prompt`, leave a placeholder for us to later insert `chat_history`, and provide a place for us to insert the user `input`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "AW96-NkUVN5u"
      },
      "outputs": [],
      "source": [
        "system_prompt = \"\"\"You are the oracle, the great AI decision maker.\n",
        "Given the user's query you must decide what to do with it based on the\n",
        "list of tools provided to you.\n",
        "\n",
        "If you see that a tool has been used (in the scratchpad) with a particular\n",
        "query, do NOT use that same tool with the same query again. Also, do NOT use\n",
        "any tool more than twice (ie, if the tool appears in the scratchpad twice, do\n",
        "not use it again).\n",
        "\n",
        "You should aim to collect information from a diverse range of sources before\n",
        "providing the answer to the user. Once you have collected plenty of information\n",
        "to answer the user's question (stored in the scratchpad) use the final_answer\n",
        "tool.\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9L59sEakdMd2"
      },
      "source": [
        "The oracle agent will be provided the tools we previously built."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "sELUig_VdXK7"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.tools import FunctionTool\n",
        "\n",
        "tools = [\n",
        "    FunctionTool.from_defaults(async_fn=fetch_arxiv),\n",
        "    FunctionTool.from_defaults(async_fn=rag_search_filter),\n",
        "    FunctionTool.from_defaults(async_fn=rag_search),\n",
        "    FunctionTool.from_defaults(async_fn=final_answer),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "id": "LNiTGVrMdtN4"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from llama_index.llms.openai import OpenAI\n",
        "\n",
        "llm = OpenAI(\n",
        "    model=\"gpt-4o\",\n",
        "    additional_kwargs={\"tool_choice\": \"required\"}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kcXeJrFZa9eG"
      },
      "source": [
        "## Events"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RzCsuWEsbAuS"
      },
      "source": [
        "We need to create a few events for our workflow. Llama-index comes with a few predefined event types, two of which we will use (`StartEvent` and `StopEvent`). However, we need to define a few additional custom event types - these are:\n",
        "\n",
        "* `InputEvent` to handle new messages and prepare chat history.\n",
        "\n",
        "* `ToolCallEvent` to trigger tool calls."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "id": "FXdBlptmbAQk"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.llms import ChatMessage\n",
        "from llama_index.core.tools import ToolSelection, ToolOutput\n",
        "from llama_index.core.workflow import Event\n",
        "\n",
        "\n",
        "class InputEvent(Event):\n",
        "    input: list[ChatMessage]\n",
        "\n",
        "\n",
        "class ToolCallEvent(Event):\n",
        "    id: str\n",
        "    name: str\n",
        "    params: dict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k9aE7MiznZCj"
      },
      "source": [
        "Now we build the workflow. Workflows consist of a single `Workflow` class with multiple `steps`. Each step is like a compute/execution step in our agentic flow.\n",
        "\n",
        "We control which step is triggered by using different `Event` types. Each step consumes a different type of event, like `InputEvent` or `ToolCallEvent`. Additionally, our workflow begins and ends with the `StartEvent` and `StopEvent` events respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "qXq1QbvScZeZ"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step\n",
        "from llama_index.core.memory import ChatMemoryBuffer\n",
        "from llama_index.core.llms import MessageRole\n",
        "\n",
        "class ResearchAgent(Workflow):\n",
        "    def __init__(\n",
        "        self,\n",
        "        *args: any,\n",
        "        oracle: OpenAI,\n",
        "        tools: list[FunctionTool],\n",
        "        timeout: int = 20,\n",
        "    ):\n",
        "        super().__init__(*args)\n",
        "        self._timeout = timeout\n",
        "        self.oracle = oracle\n",
        "        self.tools = tools\n",
        "        self.get_tool = {tool.metadata.get_name(): tool for tool in self.tools}\n",
        "        # initialize chat history/memory with system prompt\n",
        "        self.sys_msg = ChatMessage(\n",
        "            role=MessageRole.SYSTEM,\n",
        "            content=system_prompt\n",
        "        )\n",
        "        self.memory = ChatMemoryBuffer.from_defaults(llm=OpenAI())\n",
        "\n",
        "    @step\n",
        "    async def prepare_chat_history(self, ev: StartEvent) -> InputEvent:\n",
        "        # clear memory\n",
        "        self.memory = ChatMemoryBuffer.from_defaults(llm=OpenAI())\n",
        "        self.memory.put(message=self.sys_msg)\n",
        "        # get user input\n",
        "        user_input = ev.input\n",
        "        user_msg = ChatMessage(role=\"user\", content=user_input)\n",
        "        self.memory.put(message=user_msg)\n",
        "        # get chat history\n",
        "        chat_history = self.memory.get()\n",
        "        # return input event\n",
        "        return InputEvent(input=chat_history)\n",
        "\n",
        "    @step\n",
        "    async def handle_llm_input(self, ev: InputEvent) -> ToolCallEvent | StopEvent:\n",
        "        chat_history = ev.input\n",
        "        # get oracle response\n",
        "        response = await self.oracle.achat_with_tools(\n",
        "            tools=self.tools,\n",
        "            chat_history=chat_history,\n",
        "            tool_choice=\"required\",\n",
        "        )\n",
        "        # add response to chat history / memory\n",
        "        self.memory.put(message=response.message)\n",
        "        # get tool calls\n",
        "        tool_calls = self.oracle.get_tool_calls_from_response(\n",
        "            response\n",
        "        )\n",
        "        # if final_answer tool used we return to the user with the StopEvent\n",
        "        if tool_calls[-1].tool_name == \"final_answer\":\n",
        "            return StopEvent(result={\"response\": tool_calls[-1].tool_kwargs})\n",
        "        else:\n",
        "            # return tool call event\n",
        "            return ToolCallEvent(\n",
        "                id=tool_calls[-1].tool_id,\n",
        "                name=tool_calls[-1].tool_name,\n",
        "                params=tool_calls[-1].tool_kwargs,\n",
        "            )\n",
        "\n",
        "    @step\n",
        "    async def run_tool(self, ev: ToolCallEvent) -> InputEvent:\n",
        "        tool_name = ev.name\n",
        "        additional_kwargs = {\n",
        "            \"tool_call_id\": ev.id,\n",
        "            \"name\": tool_name\n",
        "        }\n",
        "        # get chosen tool\n",
        "        tool = self.get_tool.get(tool_name)\n",
        "        if not tool:\n",
        "            tool_msg = ChatMessage(\n",
        "                role=\"tool\",\n",
        "                content=f\"Tool {tool_name} not found\",\n",
        "                additional_kwargs=additional_kwargs\n",
        "            )\n",
        "        else:\n",
        "            # now call tool\n",
        "            tool_output = await tool.acall(**ev.params)\n",
        "            tool_msg = ChatMessage(\n",
        "                role=\"tool\",\n",
        "                content=tool_output.content,\n",
        "                additional_kwargs=additional_kwargs\n",
        "            )\n",
        "        self.memory.put(message=tool_msg)\n",
        "        chat_history = self.memory.get()\n",
        "        return InputEvent(input=chat_history)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pXs5pRIg4El9",
        "outputId": "f38c3fb6-4761-473c-eda1-74248ce823ff"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "research_agent.html\n"
          ]
        }
      ],
      "source": [
        "from llama_index.utils.workflow import draw_all_possible_flows\n",
        "\n",
        "draw_all_possible_flows(ResearchAgent, filename=\"research_agent.html\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cO99bN2pnLOB"
      },
      "source": [
        "Initialize the workflow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "id": "gVYXBk6SnEYE"
      },
      "outputs": [],
      "source": [
        "agent = ResearchAgent(\n",
        "    oracle=llm,\n",
        "    tools=tools,\n",
        "    timeout=60,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2ynhR2BTqRim",
        "outputId": "55778651-84eb-40ec-824f-7883ebe43724"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search_filter\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'introduction': 'Artificial Intelligence (AI) is a broad and dynamic field that encompasses the development of computer systems capable of performing tasks that typically require human intelligence. These tasks include learning, reasoning, problem-solving, perception, language understanding, and interaction. AI has evolved significantly over the years, leading to the creation of sophisticated models and applications that impact various domains.',\n",
              " 'research_steps': '1. Conducted a general search on AI to gather initial information.\\n2. Retrieved abstracts and detailed content from specific ArXiv papers related to AI.\\n3. Analyzed the content to extract relevant information about AI and its applications.',\n",
              " 'main_body': \"Artificial Intelligence (AI) has its roots in the mid-20th century, with the term 'artificial intelligence' being coined by John McCarthy in 1956. Since then, AI has grown to encompass a wide range of subfields, including machine learning, natural language processing, robotics, and computer vision. The primary goal of AI is to create systems that can perform tasks that would normally require human intelligence. These systems are designed to learn from data, adapt to new inputs, and perform human-like tasks with increasing accuracy and efficiency.\\n\\nOne of the significant advancements in AI is the development of Large Language Models (LLMs) such as GPT-4. These models are capable of understanding and generating human-like text, making them useful for a variety of applications, including chatbots, coding assistants, and design tools. For instance, the chatbot Pi, developed by Inflection, is known for its high emotional intelligence and ability to provide emotional companionship, in addition to answering questions and performing searches.\\n\\nAI agents have also made significant strides in scientific research. ChemCrow, an LLM chemistry agent, is designed to accomplish tasks in organic synthesis, drug discovery, and materials design. It integrates multiple chemistry tools and uses GPT-4 to provide specific instructions, making it a valuable assistant to expert chemists and lowering the entry barrier for non-experts.\\n\\nThe application of AI extends to various domains, including gaming, where AI agents like AlphaGo and OpenAI Five have demonstrated exceptional capabilities in playing complex games. In design, AI-powered platforms like Diagram help transform user ideas into fully editable UI designs, enhancing productivity and creativity. The versatility and potential of AI continue to expand, driven by ongoing research and development in the field.\",\n",
              " 'conclusion': 'AI is a transformative technology with the potential to revolutionize various aspects of human life. From enhancing productivity and creativity to advancing scientific research and providing emotional companionship, AI systems are becoming increasingly integrated into our daily lives. As research and development continue, the capabilities and applications of AI are expected to grow, offering new opportunities and challenges.',\n",
              " 'sources': '- An In-depth Survey of Large Language Model-based Artificial Intelligence Agents, ArXiv ID: 2309.14365\\n- Cognitive Architectures for Language Agents, ArXiv ID: 2309.02427'}"
            ]
          },
          "execution_count": 62,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "res = await agent.run(input=\"tell me about AI\")\n",
        "res[\"response\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Looyvi2I4gvc"
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBT8Wtgu6ax9"
      },
      "source": [
        "Let's test with async!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r0_Z4wwK3PbH",
        "outputId": "f06c6c5b-2090-494e-d2aa-957aff063408"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search_filter\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'introduction': 'The MAMBA model and the Transformer architecture are two significant advancements in the field of artificial intelligence, particularly in sequence modeling and natural language processing. Both have their unique characteristics and applications, making them suitable for different tasks.',\n",
              " 'research_steps': '1. Conducted a search for information on MAMBA models in AI.\\n2. Retrieved and reviewed the abstract and details of the MAMBA model from ArXiv.\\n3. Conducted a search for information on Transformer architecture in AI.\\n4. Retrieved and reviewed the abstract and details of the Transformer architecture from ArXiv.\\n5. Compiled and compared the information from both models.',\n",
              " 'main_body': \"The MAMBA model, short for 'Linear-Time Sequence Modeling with Selective State Spaces,' is a recent development in AI that aims to achieve Transformer-quality performance in a more efficient manner. According to the research, MAMBA models have shown to exceed the performance of various baselines, including modern Transformer training recipes based on LLaMa. One of the key advantages of MAMBA is its generation throughput, which is five times higher than that of Transformers of similar size. Additionally, MAMBA-3B's quality matches that of Transformers twice its size, making it a highly efficient model for language modeling tasks. The MAMBA model also outperforms state-of-the-art GAN- and diffusion-based models in terms of fidelity metrics, demonstrating its robustness in handling complex datasets.\\n\\nOn the other hand, the Transformer architecture, introduced by Vaswani et al., has been a cornerstone in the field of natural language processing. Transformers utilize a mechanism known as self-attention, which allows the model to weigh the importance of different words in a sentence when making predictions. This architecture has been the foundation for many advanced models like GPT-3 and BERT. Transformers are known for their scalability and have been successfully applied to a wide range of tasks, from language translation to text generation. However, one of the limitations of Transformers is their computational complexity, which increases quadratically with the sequence length.\\n\\nWhen comparing the two, MAMBA models offer a significant advantage in terms of computational efficiency. The linear-time complexity of MAMBA makes it more suitable for tasks that require real-time processing or involve long sequences. In contrast, Transformers, with their self-attention mechanism, provide a more nuanced understanding of the context within a sequence, making them ideal for tasks that require a deep understanding of the text. Both models have their strengths and are likely to coexist, each serving different needs in the AI landscape.\",\n",
              " 'conclusion': 'In summary, while the MAMBA model offers a more computationally efficient alternative to the Transformer architecture, the latter remains unparalleled in its ability to understand and generate complex text. The choice between the two will largely depend on the specific requirements of the task at hand, such as the need for real-time processing or the depth of contextual understanding.',\n",
              " 'sources': '- ArXiv paper on MAMBA: https://arxiv.org/abs/2312.00752\\n- ArXiv paper on Cognitive Architectures for Language Agents: https://arxiv.org/abs/2309.02427'}"
            ]
          },
          "execution_count": 63,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent1 = ResearchAgent(oracle=llm, tools=tools)\n",
        "input = \"Tell me about MAMBA models and compare them to the Transformer architecture\"\n",
        "\n",
        "\n",
        "res = await agent.run(input=input)\n",
        "res[\"response\"]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SgtozaX0Kz8J"
      },
      "source": [
        "## Clean up\n",
        "\n",
        "\n",
        "Run the below line if you'd like to delete your index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PHNFZbAVKzGw"
      },
      "outputs": [],
      "source": [
        "#pc.delete_index(index_name)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HYIM7CBsMjJm"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}